from solvers.inherit_base import InheritSolverBase
from typing import Callable, List, Tuple, Union, Iterable
import random
import numpy as np


class Individual:
    """
    每个个体都是一个服从正态分布的随机变量，出生的后代依概率出生
    """

    def __init__(self, val: float = 1, sigma: float = 0.2, age: int = 0) -> None:
        """
        Arguments:
         - `val`: 个体的值
         - `sigma`: 个体繁殖后代时的标准差
         - `age`: 个体的年龄，用于在过滤时判断个体是否应该死亡
        """
        self.val = val
        self.sigma = sigma
        self.age = age
        self._score = -np.inf

    def __call__(self, val: float = 1, sigma: float = 0.2) -> "Individual":
        return Individual(val, sigma)

    def _multiply(
        self, age_gate: int = 20, offstring_num: int = 8
    ) -> List["Individual"]:
        if random.random() > np.exp(age_gate - self.age):
            # 个体死亡
            return []
        else:
            self.age += 1
            return [
                self,
                *list(
                    Individual(abs(val), self.sigma)._change()
                    for val in np.random.normal(self.val, self.sigma, offstring_num)
                ),
            ]

    def _change(
        self,
        gate: List[float] = [0.03, 0.03, 0.1, 0.1],
        rate: List[float] = [1.25, 1.25],
    ) -> object:
        """
        用于变异的函数，每个参数都有一定的概率变异，变异的方式是乘以一个确定的数
        """
        if random.random() < gate[0]:
            self.val *= rate[0]
        if random.random() < gate[1]:
            self.val /= rate[0]
        elif random.random() < gate[2]:
            self.sigma *= rate[1]
        elif random.random() < gate[3]:
            self.sigma /= rate[1]
        return self

    def score(
        self, dft: np.ndarray, A_K: List[np.ndarray], K: int, beta: float
    ) -> float:
        if self._score != -np.inf:
            return self._score

        def _quantify(A: np.ndarray):
            B = np.zeros_like(A, dtype=int)
            A_real = np.real(A)
            A_imag = np.imag(A)
            B = (
                np.sign(A_real)
                * 2 ** np.floor((np.clip(np.log2(np.abs(A_real)), 0, None)))
                + np.sign(A_imag)
                * 2 ** np.floor((np.clip(np.log2(np.abs(A_imag)), 0, None)))
                * 1j
            )
            return B

        self._score = (
            -np.sqrt(
                np.sum(
                    np.abs(
                        beta * dft
                        - self.matmul(
                            map(lambda x: _quantify(
                                self.val * x) / self.val, A_K), K
                        )
                    )
                )
            )
            / A_K[0].shape[0]
        )
        return self._score

    def matmul(self, A: Iterable[np.ndarray], k: int):
        i = 1
        for a in A:
            if i == 1:
                temp = a
                i += 1
            else:
                temp = temp @ a
        return temp


class AlphaInherit(InheritSolverBase):
    def __init__(
        self,
        individual_num: int,
        individual: Callable[..., Individual],
        dft: np.ndarray,
        A_K: List[np.ndarray],
        K: int,
        beta: float,
        gate: float,
        val: Iterable[float],
        sigma: Iterable[float],
        lr: float,
        **kwargs,
    ) -> None:
        self.dft = dft
        self.A_K = A_K
        self.K = K
        self.beta = beta
        self.gate = gate
        self.lr = lr
        self.group = [individual(v, s) for v, s in zip(val, sigma)]

    def _filter(
        self,
        gate: float = -500,
    ) -> float:
        gate = self.gate if self.gate else gate
        score = [
            (i, individual.score(self.dft, self.A_K, self.K, self.beta))
            for i, individual in enumerate(self.group)
        ]
        score = sorted(score, key=lambda x: x[1])
        # 淘汰小于gate的个体，如果都不小于gate，则淘汰后80%的个体
        # 如果都小于，则保留后20%的个体
        for i in range(len(score)):
            if score[i][1] > gate:
                # if score[i][1] > gate and i > len(score) * 0.8:
                self.group = [self.group[j] for j, s in score[i:]]
                print(f"score: {score[-1][1]}")
                self.gate = max(score[-1][1], self.gate + self.lr)
                return score[-1][1]
        self.group = [self.group[j] for j, s in score[int(len(score) * 0.85):]]
        print(f"score: {score[-1][1]}")
        self.gate = max(score[-1][1], self.gate + self.lr)
        return score[-1][1]

    def _survive(
        self,
    ) -> object:
        score = [
            (i, individual.score(self.dft, self.A_K, self.K, self.beta))
            for i, individual in enumerate(self.group)
        ]
        score = sorted(score, key=lambda x: x[1])
        # 淘汰后80%的个体
        self.group = [self.group[i] for i, s in score[int(len(score) * 0.85):]]

    def _multiply(
        self, individual: Individual, age_gate: int = 20, **kwargs
    ) -> List[Individual]:
        return individual._multiply(age_gate=age_gate, **kwargs)
